import sqlite3
from prettytable import PrettyTable


# 初始化数据库和表
def init_db():
    conn = sqlite3.connect("space_projects.db")
    c = conn.cursor()
    # 创建项目表
    c.execute(
        """CREATE TABLE IF NOT EXISTS projects
                 (id INTEGER PRIMARY KEY, 
                  name TEXT NOT NULL,
                  description TEXT)"""
    )

    # 如果表为空，插入初始数据
    c.execute("SELECT COUNT(*) FROM projects")
    if c.fetchone()[0] == 0:
        c.execute(
            "INSERT INTO projects (name, description) VALUES (?, ?)",
            ("火星探测计划", "探索火星表面和大气层"),
        )
        c.execute(
            "INSERT INTO projects (name, description) VALUES (?, ?)",
            ("月球基地建设", "在月球建立永久科研基地"),
        )
    conn.commit()
    conn.close()


# 显示菜单
def display_menu():
    print("\n航天项目信息系统")
    print("1. 显示所有项目")
    print("2. 按关键字搜索项目")
    print("3. 添加新项目")
    print("4. 删除项目")
    print("5. 退出")
    return input("请输入选项(1-5): ")


# 显示所有项目
def display_all_projects():
    conn = sqlite3.connect("space_projects.db")
    c = conn.cursor()
    c.execute("SELECT * FROM projects")
    rows = c.fetchall()
    table = PrettyTable()
    table.field_names = ["ID", "名称", "简介"]
    for row in rows:
        table.add_row([row[0], row[1], row[2]])
    print("\n所有航天项目:")
    print(table)
    conn.close()


# 按关键字搜索项目
def search_projects(keyword):
    conn = sqlite3.connect("space_projects.db")
    c = conn.cursor()
    c.execute(
        "SELECT * FROM projects WHERE name LIKE ? OR description LIKE ?",
        (f"%{keyword}%", f"%{keyword}%"),
    )
    rows = c.fetchall()
    print(f"\n搜索结果 (关键字: {keyword}):")
    print("-" * 40)
    for row in rows:
        print(f"ID: {row[0]}")
        print(f"名称: {row[1]}")
        print(f"简介: {row[2]}")
        print("-" * 40)
    conn.close()


# 添加新项目
def add_project():
    project_id = input("请输入项目ID: ")
    name = input("请输入项目名称: ")
    description = input("请输入项目简介: ")

    conn = sqlite3.connect("space_projects.db")
    c = conn.cursor()

    # 检查ID是否存在
    c.execute("SELECT COUNT(*) FROM projects WHERE id = ?", (project_id,))
    if c.fetchone()[0] > 0:
        print("错误：该ID已存在！")
        return

    try:
        c.execute(
            "INSERT INTO projects (id, name, description) VALUES (?, ?, ?)",
            (project_id, name, description),
        )
        conn.commit()
        print("项目添加成功！")
    except sqlite3.IntegrityError as e:
        print(f"添加失败: {e}")
    finally:
        conn.close()


# 删除项目
def delete_project():
    project_id = input("请输入要删除的项目ID: ")

    conn = sqlite3.connect("space_projects.db")
    c = conn.cursor()

    c.execute("SELECT COUNT(*) FROM projects WHERE id = ?", (project_id,))
    if c.fetchone()[0] == 0:
        print("错误：该项目不存在！")
        return

    c.execute("DELETE FROM projects WHERE id = ?", (project_id,))
    conn.commit()
    print("项目删除成功！")
    conn.close()


def main():
    init_db()
    while True:
        choice = display_menu()

        if choice == "1":
            display_all_projects()
        elif choice == "2":
            keyword = input("请输入搜索关键字: ")
            search_projects(keyword)
        elif choice == "3":
            add_project()
        elif choice == "4":
            delete_project()
        elif choice == "5":
            print("感谢使用航天项目信息系统，再见！")
            break
        else:
            print("无效的选项，请重新输入！")


if __name__ == "__main__":
    main()
